#include "socks5def.h"
#include <unistd.h>
#include <iostream>
using std::cout;
using std::cerr;
using std::endl;
#include <string.h>
#include <functional>
#include <sockutil.h>
#include <arpa/inet.h>
#include <poll.h>
#include "auth_type.h"
#include "socks5_thread.h"


/**
 * @brief CONNECT_Command_errorno_to_socks5_error_code
 * X'00' 成功
 * X'01' 普通SOCKS服务器故障
 * X'02' 连接不被规则集允许
 * X'03' 网络不可达
 * X'04' 主机不可达
 * X'05' 连接被拒绝
 * X'06' TTL 过期
 * X'07' 不支持的命令
 * X'08' 不支持的地址类型
 * X'09' 到 X'FF' 未分配
 * @param err errno
 * @return
 */
CONNECTErrorCode CONNECT_Command_errorno_to_socks5_error_code(int err)
{
    switch(err) {
        case ETIMEDOUT:
            return EC_TTL_EXPIRED;
        case EPROTOTYPE:
        case EPROTONOSUPPORT:
        case EAFNOSUPPORT:
            return EC_ADDRESSTYPE_NOT_SUPPORTED;
        case ECONNREFUSED:
            return EC_CONN_REFUSED;
        case ENETDOWN:
        case ENETUNREACH:
            return EC_NET_UNREACHABLE;
        case EHOSTUNREACH:
            return EC_HOST_UNREACHABLE;
        case EBADF:
        default:
        perror("socket/connect");
        return EC_GENERAL_FAILURE;
    }
}


void CSocks5Thread::thread_worker()
{
    cout << "socks app start to use socket: " << accepted_fd_ << endl;
    char buf[1400] = {0}; //TCP DTU MSS 1400
    SOCKS_STEP cur_step = ess_connected;
    int error_code = ese_success;
    while (1)
    {
        int nrecv = recv(accepted_fd_, buf, sizeof(buf), 0);
        if (nrecv <= 0){
            error_code = ese_socket_recv_error;
            break;
        }

        switch(cur_step)
        {
        case ess_connected:
        {
            /*
             * 版本标识符/方法选择消息（version identifier/method selection message）
             * 对于此版本的协议，VER字段被固定设置为 X'05'。 NMETHODS字段表示在METHODS字段中的方法标识符（以字节为单位）的数量。
             * 当前METHOD定义的值是：
             * X'00' 无需认证（NO AUTHENTICATION REQUIRED）
             * X'01' GSSAPI
             * X'02' 用户名/密码（USERNAME/PASSWORD）
             * X'03' 至 X'7F' IANA 分配（IANA ASSIGNED）
             * X'80' 到 X'FE' 保留用于私人方法（RESERVED FOR PRIVATE METHODS）
             * X'FF' 没有可接受的方法（NO ACCEPTABLE METHODS）
            +----+----------+----------+
            |VER | NMETHODS | METHODS  |
            +----+----------+----------+
            | 1  | 1        | 1 to 255 |
            +----+----------+----------+
            */

            if (nrecv < 3)
            {
                error_code = ese_buf_too_small;
                break;
            }

            unsigned char byteVer = buf[0];
            if (byteVer != 5)
            {
                error_code = ese_wrong_ver;
            }

            unsigned char nMethods = buf[1];
            if (nrecv != 1 + 1 + nMethods)
            {
                error_code = ese_wrong_length;
            }

            /* Support ms_no_authentiacion_required & ms_password only. */
            bool method_found = false;
            for (unsigned char i=0; i!=nMethods; i++)
            {
                if (CAuthType::getInstance().getAuthorizationMode() == EAM_None && ms_no_authentiacion_required == buf[2+i])
                {
                    method_found = true;
                    if (!send_method_selection(ms_no_authentiacion_required))
                    {
                        error_code = ese_send_authorization_method_failed;
                    } else
                        cur_step = ess_authed;
                    break;
                }

                if (CAuthType::getInstance().getAuthorizationMode() == EAM_Password && ms_password == buf[2+i])
                {
                    method_found = true;
                    if (!send_method_selection(ms_password)) {
                        error_code = ese_send_authorization_method_failed;
                    } else {
                        cur_step = ess_authentica;
                        //auth_passed = true;
                    }
                    break;
                }
            }

            if (!method_found)
            {
                if (!send_method_selection(ms_no_acceptble_methods))
                {
                    error_code = ese_send_authorization_method_failed;
                }
                else
                    error_code = ese_no_authorization_method_avalible;
                break;
            }

            break;
        }
        case ess_authentica:
        {
            //
            //检查用户名密码（未来还要检查白名单等）。
            //+-------+--------------+--------+--------------+--------+
            //|Version|UserNameLength|UserName|PasswordLength|Password|
            //+-------+--------------+--------+--------------+--------+
            //|1      |      1       |  1~255 |      1       |  1~255 |
            //+-------+--------------+--------+--------------+--------+
            //Note: Version鉴权版本号,与协议版本号无关。(It is always 1 currently.)
            //Server reply:
            //+-------+------+
            //|Version|Status|
            //+-------+------+
            //|   1   |   1  |
            //+-------+------+
            //Version: same to request.(It is always 1 currently.)
            //Status: 0 for OK, other value indicate a failure code.
            const int length_minor = 5;
            if (nrecv < length_minor)
            {
                error_code = ese_buf_too_small;
                send_auth_result(1);
                break;
            }

            const char authorization_ver = 1;
            if (buf[0] != authorization_ver){
                error_code = ese_wrong_auth_ver;
                send_auth_result(1);
                break;
            }

            unsigned int len_of_user_name = buf[1];
            if (len_of_user_name + 1 + 1 >= nrecv){
                error_code = ese_invalid_data;
                send_auth_result(1);
                break;
            }

            std::string user_name;
            user_name.append((const char*)&buf[2], len_of_user_name);

            unsigned int len_of_password = buf[1+1+len_of_user_name];
            if (nrecv != 3 + len_of_user_name + len_of_password)
            {
                error_code = ese_invalid_data;
                send_auth_result(1);
                break;
            }

            std::string password;
            password.append((const char*)&buf[3+len_of_user_name], len_of_password);

            if (CAuthType::getInstance().getUserName() != user_name || CAuthType::getInstance().getPassword() != password)
            {
                error_code = ese_authorize_failed;
                send_auth_result(1);
                break;
            }

            send_auth_result(0);
            cur_step = ess_authed;

            break;
        }
        case ess_authed:
        {
            /*
             Socks5 REQUEST.
            +----+-----+-------+------+----------+----------+
            |VER | CMD | RSV   | ATYP | DST.ADDR | DST.PORT |
            +----+-----+-------+------+----------+----------+
            | 1  | 1   | X'00' | 1    | Variable | 2        |
            +----+-----+-------+------+----------+----------+
            其中：
            VER 协议版本号: X'05'
            CMD
                CONNECT X'01'
                BIND X'02'
                UDP ASSOCIATE X'03'
            RSV 保留未用（RESERVED）
            ATYP 后续地址（即DST.ADDR）的地址类型（address type of following address）
                IP V4 address: X'01'
                DOMAINNAME: X'03'
                IP V6 address: X'04'
            DST.ADDR 期望的目标地址
            DST.PORT 期望的目标端口，网络字节序
            SOCKS服务器通常会根据源地址和目标地址评估请求，并根据请求类型返回一条或多条回复消息。

            Reply:
            +----+-----+-------+------+----------+----------+
            |VER | REP | RSV   | ATYP | BND.ADDR | BND.PORT |
            +----+-----+-------+------+----------+----------+
            | 1  | 1   | X'00' | 1    | Variable | 2        |
            +----+-----+-------+------+----------+----------+
            VER 协议版本：X'05'
            REP 回复字段：
                X'00' 成功
                X'01' 普通SOCKS服务器故障
                X'02' 连接不被规则集允许
                X'03' 网络不可达
                X'04' 主机不可达
                X'05' 连接被拒绝
                X'06' TTL 过期
                X'07' 不支持的命令
                X'08' 不支持的地址类型
                X'09' 到 X'FF' 未分配
            RSV 保留未用，总是为0
            ATYP 后续地址(即BND.ADDR)的地址类型
                IP V4 地址：X'01'
                域名：X'03'
                IP V6 地址：X'04'
            BND.ADDR 服务器绑定的地址
            BND.PORT 服务器绑定的端口（网络字节序
            */

            if (nrecv < 6)
            {
                reply_connect_error(EC_GENERAL_FAILURE);
                error_code = ese_wrong_length;
                break;
            }

            unsigned char ver = buf[0];
            if (ver != 5){
                reply_connect_error(EC_GENERAL_FAILURE);
                error_code = ese_wrong_ver;
                break;
            }

            unsigned char MsgCMD = buf[1];
#define CMD_CONNECT 1
#define CMD_BIND 2
#define CMD_UDP_ASSOCIATE 3
            if (CMD_CONNECT != MsgCMD){
                //CMD_BIND & CMD_UDP_ASSOCIATE does not support yet.
                reply_connect_error(EC_GENERAL_FAILURE);
                error_code = ese_cmd_not_support;
                break;
            }

            unsigned char reserved_byte = buf[2];
            if (0 != reserved_byte){
                reply_connect_error(EC_GENERAL_FAILURE);
                error_code = ese_invalid_data;
                break;
            }
#define ADDRESS_TYPE_IPv4 1
#define ADDRESS_TYPE_DOMAIN_NAME 3
#define ADDRESS_TYPE_IPv6 4
            unsigned char address_type = buf[3];
            /*
             在地址字段（DST.ADDR、BND.ADDR）中，ATYP 字段指定字段中包含的地址类型：
             X'01' 该地址是IPv4地址，长度为4个字节。
             X'03' 地址字段包含一个完全限定域名(Full qulified domain name)。 地址字段的第一个字节为其后域名的长度（以字节为单位），长度不包括终止符NULL。
             X'04' 该地址是IPv6地址，长度为16个字节。
             */
            int addr_len = 0;
            unsigned int addr_IPv4 = 0;
            unsigned char addr_IPv6[16] = {0};
            // The port is always the last two bytes (with network byte order)
            unsigned short port = ntohs(*(unsigned short*)(&buf[nrecv-2]));

            // 存放IPv4地址或者IPv6地址的字符串形式(via inet_ntop())，或者是域名。
            // Internet presentation address for inet_ntop() function call which support IPv6 & IPv4 at the same time.
            char inet_presentation_address[256];
            memset(inet_presentation_address, 0, sizeof(inet_presentation_address));
            if (ADDRESS_TYPE_IPv4 == address_type){
                addr_len = 4;
                if (nrecv != addr_len + 6){
                    reply_connect_error(EC_GENERAL_FAILURE);
                    error_code = ese_invalid_data;
                }
                addr_IPv4 = *(unsigned int*)(&buf[4]);
                /*
                 const char *inet_ntop(int af, const void *src, char *dst, socklen_t size);
                 这个函数是随IPv6出现的函数，对于IPv4地址和IPv6地址都适用，函数中p和n分别代表表达（presentation)和数值（numeric)。
                 地址的表达格式通常是ASCII字符串，数值格式则是存放到套接字地址结构的二进制值。
                 On success, inet_ntop() returns a non-null pointer  to  dst.   NULL  is
                 returned if there was an error, with errno set to indicate the error.
                 */
                const char *ntop_result = inet_ntop(AF_INET,
                                                    &addr_IPv4,
                                                    inet_presentation_address,
                                                    sizeof(inet_presentation_address));
                if (ntop_result != inet_presentation_address)
                {
                    error_code = ese_invalid_data;
                    cerr << "inet_ntop failed with errno: " << errno << endl;
                    reply_connect_error(EC_GENERAL_FAILURE);
                }
                cout << "CONNECT CMD: target address is " << inet_presentation_address << endl;
            } else if (ADDRESS_TYPE_IPv6 == address_type) {
                addr_len = 16;
                if (nrecv != addr_len + 6){
                    reply_connect_error(EC_GENERAL_FAILURE);
                    error_code = ese_invalid_data;
                }
                memcpy(addr_IPv6, &buf[4], 16);
                const char* ntop_result_v6 = inet_ntop(AF_INET6,
                                                       &addr_IPv6,
                                                       inet_presentation_address,
                                                       sizeof(inet_presentation_address));
                if (ntop_result_v6 != inet_presentation_address){
                    error_code = ese_invalid_data;
                    cerr << "inet_ntop failed with errno: " << errno << endl;
                    reply_connect_error(EC_GENERAL_FAILURE);
                }
                cout << "CONNECT CMD: target address is: " << inet_presentation_address << endl;
            } else if (ADDRESS_TYPE_DOMAIN_NAME == address_type){
                addr_len = (unsigned char)buf[4];
                if (nrecv != addr_len + 7){
                    reply_connect_error(EC_GENERAL_FAILURE);
                    error_code = ese_invalid_data;
                    break;
                }
                memcpy(inet_presentation_address, buf+5, addr_len);
                cout << "CONNECT CMD: target address is: " << inet_presentation_address << endl;
            } else {
                addr_len = 0;
                reply_connect_error(EC_GENERAL_FAILURE);
                error_code = ese_invalid_data;
                break;
            }

            if (0 == error_code){
                addrinfo* remote_addr = NULL;
                if (0 != resolve(inet_presentation_address, port, &remote_addr))
                {
                    cerr << "CONNECT: CANNOT resolve " << inet_presentation_address << ":" << port << endl;
                    error_code = ese_invalid_data;
                    reply_connect_error(EC_GENERAL_FAILURE);
                } else {
                    int sock_remote_fd = socket(remote_addr->ai_addr->sa_family, SOCK_STREAM, 0);
                    if (sock_remote_fd == -1)
                    {
                        cerr << "socket() failed with error: " << errno << endl;
                        reply_connect_error(CONNECT_Command_errorno_to_socks5_error_code(errno));
                    } else {
                        //do connect
                        //如有需要可以在这里用bind指定本机端口，这里我们让系统自动分配。
                        if (-1 == connect(sock_remote_fd, remote_addr->ai_addr, remote_addr->ai_addrlen))
                        {
                            cerr << "connect() remote target failed with error: " << errno << endl;
                            reply_connect_error(CONNECT_Command_errorno_to_socks5_error_code(errno));
                        }
                        else
                        {
                            //已同时连接上源端fd和远端fd，我们将分别从两个fd中读取内容并将读出的内容发往另一个fd。
                            sock_fd_remote_ = sock_remote_fd;
                            cout << "forward connection built: src fd:" << accepted_fd_ << ", dest fd:" << sock_fd_remote_ << endl;
                            reply_connect_error(EC_SUCCESS);
                            forword_loop();
                        }
                        //...
                    }
                    freeaddrinfo(remote_addr);
                }
            }
            break;
        }
        default:
            break;
        }

        if (ese_success != error_code){
            cerr << "A error " << error_code << " is caught" << endl;
            break;
        }
    }
    if (accepted_fd_ > 0)
        close(accepted_fd_);
    if (sock_fd_remote_ > 0)
        close(sock_fd_remote_);
    is_thread_end_ = true;
    cout << "work thread is finished, accepted fd:" << accepted_fd_ << endl;
}

bool CSocks5Thread::send_method_selection(MethodSelection ms)
{
    /*
    +----+--------+
    |VER | METHOD |
    +----+--------+
    | 1  | 1      |
    +----+--------+
    */
    char method = (char)ms_no_acceptble_methods;
    switch (ms)
    {
    case ms_no_authentiacion_required:
    case ms_password:
    case ms_no_acceptble_methods:
        method = (char)ms;
        break;
    }

    char buf[2] = {5, 0};
    buf[1] = method;

    int ilen = send(accepted_fd_, buf, sizeof(buf), 0);
    if (ilen != sizeof(buf))
    {
        cerr << "method_selection send bytes: " << ilen << endl;
        cerr << "method_selection send data failed with error: " << errno << endl;
    }
    return ilen == sizeof(buf);

}

bool CSocks5Thread::send_auth_result(unsigned char result)
{
    char buf[2] = {1, 0};
    if (result != 0){
        buf[1] = result;
    }

    int ilen = send(accepted_fd_, buf, sizeof(buf), 0);
    if (ilen != sizeof(buf))
    {
        cerr << "send_auth_result send bytes: " << ilen << endl;
        cerr << "send_auth_result send data failed with error: " << errno << endl;
    }
    return ilen == sizeof(buf);
}

/*
+----+-----+-------+------+----------+----------+
|VER | REP | RSV   | ATYP | BND.ADDR | BND.PORT |
+----+-----+-------+------+----------+----------+
| 1  | 1   | X'00' | 1    | Variable | 2        |
+----+-----+-------+------+----------+----------+
*/
bool CSocks5Thread::reply_connect_error(CONNECTErrorCode ec)
{
    //Can I always use IPv4 address?
    if (ec != 0)
        cout << "reply_connect_error:" << ec << endl;
    char reply[10] = {5, ec, 0, 1, 0, 0, 0, 0, 0, 0};
    return sizeof(reply) == send(accepted_fd_, reply, sizeof(reply), 0);
}

void CSocks5Thread::forword_loop()
{
    //assert(accepted_fd_ != -1);
    //assert(m_sock_fd_remote_ != -1);

    //int poll(struct pollfd *fds, nfds_t nfds, int timeout);
    // 在代理双向的socket上监听数据包到来的事件
    pollfd fds[2];
    memset(fds, 0, sizeof(fds));
    fds[0].fd = accepted_fd_;
    fds[1].fd = sock_fd_remote_;
    fds[0].events = fds[1].events = POLLIN;

    // timeout: milliseconds.
    while (!is_force_stop_){
        int r = poll(fds, 2, 1000);
        if (0 == r)
        {
            //cout << "[" << accepted_fd_ << "," << sock_fd_remote_ << "] time out." << endl;
            continue;
        }
        else if (-1 == r){
            if (EINTR == errno || EAGAIN == errno)
            {
                cout << "EINTR" << endl;
                continue;
            }
            else
            {
                cerr << "poll() failed with error no: " << errno << endl;
                break;
            }
        } else {
            char buf[40960];
            int recv_fd, send_fd;
            if (fds[0].revents & POLLIN){
                recv_fd = fds[0].fd;
                send_fd = fds[1].fd;
            } else {
                recv_fd = fds[1].fd;
                send_fd = fds[0].fd;
            }

            int n = recv(recv_fd, buf, sizeof(buf), 0);
            //cout << "recv bytes: " << n << endl;
            if (n > 0){
                int sent_bytes = 0;
                while (sent_bytes < n){
                    int m = send(send_fd, buf+sent_bytes, n-sent_bytes, 0);
                    //cout << "send bytes: " << m << endl;
                    if (m < 0){
                        cerr << "send error: " << errno << endl;
                        return;
                    }
                    sent_bytes += m;
                }
            } else {
                break;
            }
        }
    }

}

CSocks5Thread::CSocks5Thread(int accepted_fd)
    : accepted_fd_(accepted_fd)
    , sock_fd_remote_(-1)
    , is_thread_end_(false)
    , is_force_stop_(false)
    , theThread_(new std::thread(std::bind(&CSocks5Thread::thread_worker, this)))
{

}

CSocks5Thread::~CSocks5Thread(){
    join();
}

bool CSocks5Thread::get_is_thread_end() const
{
    return is_thread_end_;
}

void CSocks5Thread::join(){
    if (theThread_ && theThread_->joinable())
        theThread_->join();
}

void CSocks5Thread::stop() {
    is_force_stop_ = true;
    if (accepted_fd_ > 0){
        close(accepted_fd_);
    }
    if (sock_fd_remote_ > 0)
        close(sock_fd_remote_);
}
